'''
Created on Feb 2, 2010

@author: oabalbin
'''

#import sys;
#import string;
#import StringIO;
import os;
import sys;
import glob;
import numpy as np
import scipy as sp
from scipy import stats

from optparse import OptionParser
from collections import deque

# My libraries
import signatures.parsers.read_microarray_data as pm
#import signatures.parsers.read_gene_lists as gpm
import signatures.db.tables as dt
import signatures.db.query as dq
import signatures.common.classes as cm #cm #record, arrayMatrix
import signatures.preprocess.tools as tl #filter_constant_genes, get_variable_genes, 
import signatures.plotting.plots as myplt

class pipeline():
    
    def __init__(self, dbhost, dbuser, dbpasswd, database):
        
        self.mt = dt.tables(dbhost, dbuser, dbpasswd, database)
        self.mq = dq.query(dbhost, dbuser, dbpasswd, database)
        #table information
        #NSC,CONCUNIT,LCONC,PANEL,CELL,PANELNBR,CELLNBR,NLOGGI50,INDN,TOTN,STDDEV        
        #17 ,M ,-4.00,Non-Small Cell Lung ,NCI-H23 ,1 ,1 ,4.021,2 ,2 ,0.029
        
        self.tableDescription = '''NSC_ID INT, CONCUNIT VARCHAR(40), LCONC FLOAT(20,14), PANEL VARCHAR(60), CELL VARCHAR(40), 
                                 PANELNBR SMALLINT, CELLNBR SMALLINT, NLOGGI50 FLOAT (20,14), INDN SMALLINT,TOTN SMALLINT, STDDEV FLOAT(20,14), 
                                 INDEX NSC_ID (NSC_ID), INDEX CELL (CELL), INDEX PANEL (PANEL)'''
        self.tableName = 'Drugs_GI50'
        
        # number of samples a compound should be in to be considered
        self.nmasked = 40
        self.nmeas = 20
    
    
    def parse_input_annotfiles(self, annotInputFile):
    
        expr_file, outdirs, samples, typefile={},{},{},'nofile'
    
        for line in annotInputFile:
            fields = line.strip('\n')
                 
            if fields[0][0] == '' or fields[0][0] == '#':
                continue
            
            if fields[0][0] == '@':
                typefile = fields
                continue
            
            fields = fields.split('=')
            dict = {'@Array expression':expr_file, '@Samples':samples,'@Output Directories':outdirs}[typefile]
            dict[fields[0]] = fields[1]
            
        return expr_file, samples, outdirs
    

    def read_files_folder(self,folderpath,ext):
        ''' '''
        # Read files in folder
        myfiles=[]
        for infile in glob.glob( os.path.join(folderpath, '*'+ext) ):
            myfiles.append(infile)
            #print "current file is: " + infile
        return myfiles
    
    
    def process_microarray_files(self,exp_files,tp):
        dbf=[]
        for f in exp_files:
            of=f+'_dbfile'
            tp.microarray_nci60_file(open(f),open(of,'w'))
            dbf.append(of)
        
        return dbf
    
    def create_database(self,dbfiles,doit=False):
        """
        It creates a databases if doit = True
        """
        if doit:
            print 'Making tables'            
            for dbf in dbfiles:     
                if not(self.mt.table_existence(self.tableName)):
                    self.mt.create_table(self.tableName,self.tableDescription)
                    
                self.mt.load_data_table_cvs(self.tableName,dbf)
    
    
    def write_compound_mat(self, expMat, vargenes, matfile, genefile):
        """
        write an compound dataset. 
        """
        
        genelist = np.array(expMat.geneSym)
        genelist = genelist[vargenes]
        
        vg_expMat = expMat.expVal[vargenes,:]
        
        samples = [ i[1] for i in expMat.sample]
        matfile.write(",".join(samples).replace(",", "\t")+'\n')
        for i, g in enumerate(genelist):
            el = map( str,list(vg_expMat[i,:]) ) 
            genefile.write(g+'\n')
            matfile.write(g+"\t"+",".join(list( el )).replace(",", "\t") + '\n')
    
    def pickledump_compound_mat(self, expMat, varcpm, matfile, genefile):
        """
        Pickle the compund data set, so we can used for later analysis as 
        the correlation with factor activity.
        varcpm allows to subset the data
        """
        
        genelist = np.array(expMat.geneSym)
        genelist = genelist[varcpm]
        vg_expMat = expMat.expVal[varcpm,:]
        samples = [ i[1] for i in expMat.sample]
        genes = [g for g in genelist]
        genefile.write(",".join(samples).replace(",", "\t")+'\n')
        genefile.write(",".join(genes).replace(",", "\t")+'\n')
        np.ma.dump(vg_expMat, matfile)
        
             
    
    def parse_drugmatrix_file(self,inputfile,maskemode=False):
        """
        It parses a drug sensitivity matrix: rows are drug sensitivity, columns are cell lines i which the drug sensitivity was evaluated.  
        """
        ncs_name, chem_name, GI50_vals, header = deque(),deque(),deque(), deque()
        header = inputfile.next().strip('\n').split('\t')
        samples=header[2:]
        for line in inputfile:
            fields = line.strip().split('\t')            
            if len(fields[2:]) < len(samples) or len(fields[2:]) > len(samples):
                print fields[0], fields[1]
                sys.exit(0)
            ncs_name.append(fields[0]), chem_name.append(fields[1]), GI50_vals.append(map(float,fields[2:]))
        
        
        number_of_compounds, number_of_samples = len(ncs_name),len(samples)
        expmat = cm.cpmMatrix(number_of_compounds, number_of_samples)
        [expmat.sample.append((j,sp)) for j, sp in enumerate(samples)]
        
        i = 0
        for ncs, chem, GI50 in zip(ncs_name, chem_name, GI50_vals):
            expmat.probeID.append(chem)
            expmat.geneSym.append(str(ncs))
            expmat.geneSynd.add((i,chem))            
            expmat.expValu[i,:] = np.array(GI50)
             
            if maskemode:
                mode_ind = mask_mode_GI50vals(np.array(GI50))

            thf = np.isnan(GI50)
            common_mask = np.ma.mask_or(mode_ind, thf)
            expmat.expmask[i,:]=common_mask
            
            expmat.stdVal[i,j] = np.ma.std(GI50)
            i+=1
            
        expmat.create_gen_index()
        
        return expmat


def mask_mode_GI50vals(myarray):
    """
    This function indentifies GI50 values that represent no-effect in the cell lines. The criteria is that a particular gi50 value is present 
    in many cell lines is maybe because the experiment did not work or because it is the maximum concentration of that particular compound.  
    """
    # 10% of the data
    cutoff=int(0.1*myarray.shape[0])
    
    mode_val = sp.stats.mode(myarray)
    
    if  mode_val[1] >= cutoff:
        mode_indicator = myarray == mode_val[0]
    else:
        mode_indicator = np.ma.make_mask_none((myarray.shape[0],))
        
    return mode_indicator

        

if __name__ == '__main__':
    
    optionparser = OptionParser("usage: %prog [options] ")
    optionparser.add_option("-f", "--annotFile", dest="annotFile",
                            help="annotation file for all files to use")
    optionparser.add_option("-m", action="store_true", dest="localMat",
                            help="uses a local matrix of drug sensitivities instead of database")
    optionparser.add_option("-d", action="store_true", dest="dumpcorr",
                            help="Dump the outa correlation mat. If false it loads a preexistent one")

    
    
    (options, args) = optionparser.parse_args()
       
    # start the pipeline. 
    # open the database
    pipe = pipeline("localhost", "oabalbin", "oscar", "signatures")
    expr_folder, samples, outdirs = pipe.parse_input_annotfiles(open(options.annotFile))    
    
    tp = pm.parser()
    ####### output files
    
    outfile1 = outdirs['outdir']+'drugsensitivy_mat_standardAgents.txt'  # Expression values
    outfile2 = outdirs['outdirGeneral']+'permutations_matcorr.txt'
    outfile3 = outdirs['outdir']+'pathcorr_dump_standardAgents.txt'
    ######### Get raw data  ############
    #### Get sample names
    # mat1 is the drug matrix. mat2 is the activity matrix
    # mat2 is the path activity matrix
    samplelist = tp.list_of_names_in_line(open(samples['samples']))
    mat1_hd = tl.get_header_mat(open(expr_folder['cmpsamp']))

    # load the compound matrix object and the path activity
    if options.localMat:
        mat1_raw = pipe.parse_drugmatrix_file(open(expr_folder['mat4']),True)
    else:
        mat1_raw = tl.pickleload_matrix_data(expr_folder['mat1'])
    
    # Creating a masked array with the values of oathway activity
    cols = range(61)
    mat2_unmasked = np.loadtxt(expr_folder['mat2'],dtype=float, delimiter='\t',skiprows=1, usecols=cols[1:]) #tl.get_matrix_data(open(pathactmat))
    mat2mask = np.ma.make_mask_none((mat2_unmasked.shape[0],mat2_unmasked.shape[1]))    
    mat2_raw = np.ma.array(mat2_unmasked,mask=mat2mask)
    # gene names
    gene_names = np.loadtxt(expr_folder['mat2'],dtype=str, delimiter='\t',skiprows=1, usecols=[0])
    
    # Rearrange the compound matrix to match the column order of the
    # activity matrix. This is required in order to do the correlations    
    odmat1_raw = tl.order_arrayMat_bysamples(mat1_raw, samplelist, mat1_hd)
    ########## Center and Normalize the values in the matrix
    # normalize the values for the matrix of the drug sensitivity in the NCI60:
    mat1_raw.expVal = tl.center_normalize_mat(mat1_raw.expVal, True)
    # normalize the values for the matrix of the path activities in the NCI60:
    mat2 = tl.center_normalize_mat(mat2_raw, True)

    # Rearrange the compound matrix to match the column order of the
    # activity matrix. This is required in order to do the correlations
    odmat1 = tl.order_arrayMat_bysamples(mat1_raw, samplelist, mat1_hd)  
    
      
    print samplelist
    #print odmat1.sample
    print mat2.shape
    print mat1_raw.expVal.shape
    #print mat1.expVal
        
    # pickle the correlation matrix
    if options.dumpcorr:   
        outa = tl.calc_correlation_mats_mod(odmat1.expVal, mat2, odmat1.geneSym, 'Robs', 0.05, 0.2, outfile2, True) #outfile2
                 
        #print 'Dumping the correlation array' 
        outa.dump(outfile3)
    else:
        print 'Loading the correlation array'
        outa = tl.pickleload_matrix_data(open(expr_folder['mat3']))

    print outa.shape
    # normalize the values for the matrix of the path activities in the NCI60:
    mat2_norm1 = tl.normalize_matrix_rows(mat2_raw) 
    
    # center and normalized the GI50 of correlated compounds
    # odmat1.expVal[outa[:,0],:] ==> only the compounds with good correlation
    #cmp = map(int,outa[:,0])
    #odmat1_norm1 = tl.normalize_matrix_rows(odmat1_raw.expVal)
    
    
    myplt.write_drugpath_mat(odmat1.expVal,mat2, outa, outdirs['outdir'], samplelist, True, odmat1.probeID, list(gene_names))
    # Plot the correlated compounds
    myplt.plot_drugvsact(odmat1_raw.expVal,mat2, outa, outdirs['outdir'], True, list(gene_names))
    
    sys.exit(0)